﻿using Microsoft.ServiceBus;
using Microsoft.ServiceBus.Messaging;
using Microsoft.WindowsAzure.Storage;
using Microsoft.WindowsAzure.Storage.Queue;
using Microsoft.WindowsAzure.Storage.Table;
using System;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Threading;
using System.Threading.Tasks;



namespace vmagent
{
    class vmagent
    {

        // Name of queues
        private const string HIGHQ = "highq";
        private const string MEDIUMQ = "mediumq";
        private const string LOWQ = "lowq";
        private const string IDLE = "idlequeue";
        //name of job table
        private const string TABLE = "jobtable";
        //name of command queue
        private const string CQNAME = "cqueue";
        //name of status queue
        private const string RQNAME = "outqueue";
        //name of service bus
        private const string SBNAME = "<service namespace>";
        //service bus key
        private const string SBKEY = "<service bus key>";
        //service bus owner
        private const string SBISSUER = "owner";

        // Max idle time before an instance signals that it is idle
        // 20 minutes in msec
        private const int _maxidletime = 20 * 60 * 1000;
        // Max duration of a job = visibility timeout to be set
        static TimeSpan _maxjobtime = TimeSpan.FromMinutes(20);

        // set these constants with your values to run the sample.

        private const string _storageacc = "<storage account name>";
        private const string _storagekey = "<storage key>";

        //Connection Strings
        private const string _sbconnectionstring = "Endpoint=sb://<service namespace>.servicebus.windows.net;SharedSecretIssuer=owner;SharedSecretValue=<service bus key>";
        private const string _connectionstring = "DefaultEndpointsProtocol=https;AccountName=<storage account name>;AccountKey=<storage key>";


        public class JobEntity : TableEntity
        {
            public JobEntity(string node, string jobid)
            {
                this.PartitionKey = node;
                this.RowKey = jobid;
            }
            public JobEntity() { }
            public string Status { get; set; }
            public DateTime TimeIn { get; set; }
            public DateTime TimeUpdated { get; set; }
        }
        static void Main(string[] args)
        {

            Console.WriteLine("Time: {0}, VMAgent starting.", DateTime.Now);
            CloudStorageAccount storageAccount = CloudStorageAccount.Parse(_connectionstring);
            CloudQueueClient queueclient = storageAccount.CreateCloudQueueClient();
            CloudQueue highq = queueclient.GetQueueReference(HIGHQ);
            highq.CreateIfNotExists();
            CloudQueue mediumq = queueclient.GetQueueReference(MEDIUMQ);
            mediumq.CreateIfNotExists();
            CloudQueue lowq = queueclient.GetQueueReference(LOWQ);
            lowq.CreateIfNotExists();
            CloudQueue idlequeue = queueclient.GetQueueReference(IDLE);
            idlequeue.CreateIfNotExists();

            CloudTableClient tableclient = storageAccount.CreateCloudTableClient();
            CloudTable jobtable = tableclient.GetTableReference(TABLE);
            jobtable.CreateIfNotExists();

            var sbnamespace = NamespaceManager.CreateFromConnectionString(_sbconnectionstring);

            //connect to results topic
            TopicClient rqclient = TopicClient.CreateFromConnectionString(_sbconnectionstring, RQNAME);

            // node name
            string node = Dns.GetHostName();
            // default n. of messages picked = n. of processors, but you can set to any value
            int maxmsgs = Environment.ProcessorCount;
            ParallelOptions poptions = new ParallelOptions();
            poptions.MaxDegreeOfParallelism = maxmsgs;

            bool exclusive = false;
            int running = 0;

            //execute as many threads as max messages. Each thread deals with a msg hence a task to run
            Parallel.For(0, maxmsgs, poptions, i =>
            {
                //idle time in msec
                int idletime = 0;

                do
                {
                    // read the queues in order of priority
                    var msg = highq.GetMessage(visibilityTimeout: _maxjobtime);

                    //var msg = highq.GetMessage(visibilityTimeout: _maxjobtime);
                    var taskqueue = highq;
                    if (msg == null)
                    {
                        msg = mediumq.GetMessage(visibilityTimeout: _maxjobtime);
                        taskqueue = mediumq;
                    }
                    if (msg == null)
                    {
                        msg = lowq.GetMessage(visibilityTimeout: _maxjobtime);
                        taskqueue = lowq;
                    }

                    if (msg != null)
                    {
                        idletime = 0;
                        Console.WriteLine(msg.AsString);
                        /* the format of the message will be:
                         * JobID command argument1 argument2...
                         * the separator is a space ' '*/
                        string[] content = msg.AsString.Split(' ');
                        string jobid = content[0];
                        if (content[1] == "exclusive")
                            exclusive = true;
                        else
                            exclusive = false;
                        string command = content[2];
                        string pars = "";
                        if (content.Length > 3)
                        {
                            pars = msg.AsString.Substring(msg.AsString.IndexOf(content[3]));
                        }

                        if (exclusive && (running != 0))
                        {
                            //make the message visible again if it is exclusive and this node is busy already;
                            taskqueue.UpdateMessage(msg, TimeSpan.FromSeconds(0.0), MessageUpdateFields.Visibility);
                            //pause this thread for 10 seconds to give some time for exclusive job to finish;
                            Console.WriteLine("Pausing thread 10 seconds for other tasks to finish.");
                            Thread.Sleep(10000);
                        }
                        else
                        {
                            running++;
                            Console.WriteLine("Running tasks: {0}", running);
                            Console.WriteLine("JobID: {0}, Exclusive: {3}, Command: {1}, Parameters:{2}", jobid, command, pars, exclusive);
                            var job = jobquery(jobtable, jobid);
                            if (job.Status == "Cancelled")
                            {
                                running--;
                                Console.WriteLine("JobID: {0}, Command: {1} was cancelled before execution started.", jobid, command);
                                taskqueue.DeleteMessage(msg);
                            }
                            else
                            {
                                //update the job status to running
                                update_status(jobtable, node, job, "Running", rqclient);
                                //create a subscription to the job topic and filter for this job
                                Random randomn = new Random();
                                //max 256 subscriptions per node, as a node could in theory run 256 tasks, 1 per core supported in 2008r2
                                string subname = "JobCommands-" + node + "-" + randomn.Next(1, 256) + i;
                                // set filter to jobid
                                SqlFilter jobfilter = new SqlFilter("JobID = " + "\'" + jobid + "\'");
                                if (sbnamespace.SubscriptionExists(CQNAME, subname))
                                    sbnamespace.DeleteSubscription(CQNAME, subname);
                                sbnamespace.CreateSubscription(CQNAME, subname, jobfilter);
                                //subscribe to commands and filter by jobid actually owned
                                SubscriptionClient subclient = SubscriptionClient.CreateFromConnectionString(_sbconnectionstring, CQNAME, subname);

                                //now try and spawn the executable

                                try
                                {
                                    Process p = new Process();

                                    p.StartInfo.UseShellExecute = false;
                                    p.StartInfo.CreateNoWindow = false;
                                    p.StartInfo.Arguments = pars;
                                    p.StartInfo.FileName = command;

                                    p.Start();
                                    //while the process is executing, listen to commands for it
                                    int returncode = 0;
                                    // catch cancel message but do not poll the sb 
                                    subclient.OnMessage((submsg) =>
                                    {
                                        //if (submsg != null)
                                        //{
                                        try
                                        {
                                            submsg.RenewLock();
                                            string msgbody = submsg.GetBody<string>();
                                            Console.WriteLine("Command received: {0}", msgbody);
                                            if (msgbody == "CANCEL")
                                            {
                                                p.Kill();

                                                returncode = -9;
                                                submsg.Complete();
                                            }

                                        }
                                        catch (Exception ex)
                                        {
                                            Console.WriteLine(ex.Message);
                                            submsg.Abandon();
                                        }
                                    });
                                    // }
                                    // wait for the process to terminate
                                    while (!p.HasExited) Thread.Sleep(1000);

                                    if (returncode != -9) returncode = p.ExitCode;
                                    switch (returncode)
                                    {
                                        //the job executed successfully
                                        case 0:
                                            {
                                                running--;
                                                taskqueue.DeleteMessage(msg);
                                                update_status(jobtable, node, job, "Completed", rqclient);
                                                sbnamespace.DeleteSubscription(CQNAME, subname);
                                                Console.WriteLine("JobID: {0}, Command: {1} completed.", jobid, command);
                                                break;
                                            }
                                        //the job was cancelled
                                        case -9:
                                            {
                                                running--;
                                                taskqueue.DeleteMessage(msg);
                                                update_status(jobtable, node, job, "Cancelled", rqclient);
                                                sbnamespace.DeleteSubscription(CQNAME, subname);
                                                Console.WriteLine("JobID: {0}, Command: {1} was cancelled.", jobid, command);
                                                break;
                                            }
                                        //in any other case the job is assumed to fail.
                                        default:
                                            {
                                                running--;
                                                taskqueue.DeleteMessage(msg);
                                                update_status(jobtable, node, job, "Failed", rqclient);
                                                sbnamespace.DeleteSubscription(CQNAME, subname);
                                                Console.WriteLine("JobID: {0}, Command: {1} failed and returned {2}.", jobid, command, returncode);
                                                break;
                                            }
                                    }
                                }
                                catch (Exception ex)
                                {
                                    //the process could not be started at all or had other o/s-related problems
                                    running--;
                                    taskqueue.DeleteMessage(msg);
                                    update_status(jobtable, node, job, "Failed", rqclient);
                                    sbnamespace.DeleteSubscription(CQNAME, subname);
                                    Console.WriteLine("JobID: {0}, Command: {1} failed and caused an exception.", jobid, command);
                                    Console.WriteLine(ex.Message);
                                }
                            }
                        }
                    }
                    else
                    {
                        //no msg, wait 10 seconds
                        Thread.Sleep(10000);
                        idletime += 10000;
                        Console.WriteLine("No task. Idle for {0} out of a max of {1} seconds.", idletime / 1000, _maxidletime / 1000);
                    }

                } while (idletime < _maxidletime); //all threads must be idle for node to be idle
            }); //end parallel for

            /*now I have waited too long (longer than maxidle time), signal I am idle
             * and shutdown or just stop accepting new jobs - autoscaler will remove or restart */
            CloudQueueMessage idlemessage = new CloudQueueMessage(node);
            idlequeue.AddMessage(idlemessage);
            Console.WriteLine("Time: {0}, Node: {1} going idle - signalling autoscaler.", DateTime.Now, node);
            //Comment out for debugging!
            //Process.Start("shutdown", "/s /t 60");
        }


        static JobEntity jobquery(CloudTable jobtable, string jobid)
        {

            TableQuery<JobEntity> query = new TableQuery<JobEntity>().Where(TableQuery.GenerateFilterCondition("RowKey", QueryComparisons.Equal, jobid));
            var results = jobtable.ExecuteQuery(query);

            Console.WriteLine("Node\t\tJobID\t\t\tStatus\t\tTimeIn\t\t\tTimeUpdated");

            foreach (var item in results)
            {
                Console.WriteLine("{0:-10}\t\t{1:-16}\t{2:-10}\t{3:G}\t{4:G}", item.PartitionKey, item.RowKey, item.Status, item.TimeIn, item.TimeUpdated);
            }
            JobEntity result = results.LastOrDefault<JobEntity>();

            return result;
        }


        //update the job status table
        static void update_status(CloudTable jobtable, string node, JobEntity job, string newstatus, TopicClient rqclient)
        {
            Console.WriteLine("JobID: {0}, Current Status: {1}", job.RowKey, job.Status);
            /*the node that is updating the status is not necessarily the same as the one performing the last update
             * For instance one such node could have submitted, another run the job*/
            job.PartitionKey = node;
            job.Status = newstatus;
            job.TimeUpdated = DateTime.Now;
            BrokeredMessage msg = new BrokeredMessage(newstatus);
            msg.Properties["JobID"] = job.RowKey;
            TableOperation update = TableOperation.InsertOrReplace(job);
            jobtable.Execute(update);
            rqclient.Send(msg);
            Console.WriteLine("JobID: {0}, New Status: {1}, TimeUpdated: {2}", job.RowKey, job.Status, job.TimeUpdated);

        }
    }
}
